/******************************************************************************
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * libkperf licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
 * PURPOSE.
 * See the Mulan PSL v2 for more details.
 * Author:
 * Create: 2025-05-13
 * Description: Find processes by keyword, analyze thread LLC miss ratio and bind to CPU cores
 * Current capability: Bind complementary threads to the same physical core based on LLC miss ratio
 ******************************************************************************/
#include <iostream>
#include <vector>
#include <algorithm>
#include <string>
#include <map>
#include <sstream>
#include <fstream>
#include <cstring>
#include <iomanip>
#include <unistd.h>
#include <unordered_map>
#include <unordered_set>
#include <signal.h>
#include <chrono>
#include <pthread.h>
#include <sched.h>
#include <sys/types.h>
#include <dirent.h>
#include "pcerrc.h"
#include "pmu.h"
#include "symbol.h"

static std::map<unsigned, double> numaTotalDDRC;                  //numa Id --> average ddrc bandwidth
static std::unordered_map<unsigned, unsigned*> numaToCpuCore;     //numa Id --> cpu core ids
static std::unordered_map<unsigned, unsigned> numaToCpuNumber;    //numa Id --> length of cpu cores
static std::vector<int> pidBoundCpus;                             //bounded cpus of designated pid
static unsigned numaNum = 0;                                      //number of NUMAs

const int FLOAT_PRECISION = 2;
const int TIME_UNIT_TRANS = 1000;

uint64_t topNum = 0;
uint64_t duration = 10;  // 固定为10秒统计时间
uint64_t period = 0;

// 新增：缓存访存阈值，低于此值的线程被认为不活跃
const unsigned CACHE_ACCESS_THRESHOLD = 10000;  // 10K次缓存访问作为活跃度阈值

// Store thread ID and corresponding LLC miss ratio
struct ThreadMissRatio {
    int pid;           // Process ID
    unsigned tid;      // Thread ID
    unsigned llc_miss; // LLC miss count
    unsigned llc_cache; // LLC cache count
    double ratio;      // LLC miss ratio
    std::string comm;  // Process name
    bool is_active;    // 是否活跃（缓存访问数超过阈值）

    ThreadMissRatio(int p, unsigned t, unsigned miss, unsigned cache, const std::string& process_name = "")
        : pid(p), tid(t), llc_miss(miss), llc_cache(cache), comm(process_name) {
        ratio = (cache != 0) ? static_cast<double>(miss) / cache : 0.0;
        is_active = (cache >= CACHE_ACCESS_THRESHOLD);
    }

    bool operator<(const ThreadMissRatio& other) const {
        return ratio > other.ratio; // Sort by LLC miss ratio in descending order
    }
};

// Get system CPU core count (limited to 0-79 range)
int getSystemCpuCount() {
    // Since user specified core range as 0-79, directly return 80
    return 80;
}

// Get physical core to logical cores mapping
// Returns a vector of pairs where each pair represents a physical core's two logical cores
std::vector<std::pair<int, int>> getPhysicalCoreMapping() {
    std::vector<std::pair<int, int>> physicalCores;
    // In a typical hyperthreaded system, adjacent cores (0,1), (2,3), etc. share the same physical core
    for (int i = 0; i < 80; i += 2) {
        physicalCores.push_back(std::make_pair(i, i + 1));
    }
    return physicalCores;
}

// Get even core list (even numbers in 0-79 range)
std::vector<int> getEvenCores() {
    std::vector<int> evenCores;
    for (int i = 0; i < 80; i += 2) {
        evenCores.push_back(i);
    }
    return evenCores;
}

// Get odd core list (odd numbers in 0-79 range)
std::vector<int> getOddCores() {
    std::vector<int> oddCores;
    for (int i = 1; i < 80; i += 2) {
        oddCores.push_back(i);
    }
    return oddCores;
}

// Bind thread to specified CPU core
bool bindThreadToCore(int tid, int coreId) {
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(coreId, &cpuset);

    // Use sched_setaffinity to bind thread to specified core
    if (sched_setaffinity(tid, sizeof(cpu_set_t), &cpuset) == -1) {
        std::cerr << "Failed to bind thread " << tid << " to core " << coreId << ": " << strerror(errno) << std::endl;
        return false;
    }

    return true;
}

void totalDDRCBandwidth()
{
    PmuDeviceAttr devAttr[2];
    devAttr[0].metric = PMU_DDR_READ_BW;
    devAttr[1].metric = PMU_DDR_WRITE_BW;
    int pd = PmuDeviceOpen(devAttr, 2);
    PmuEnable(pd);
    sleep(1);
    PmuData *oriData = nullptr;
    int oriLen = PmuRead(pd, &oriData);
    PmuDeviceData *devData = nullptr;
    auto len = PmuGetDevMetric(oriData, oriLen, devAttr, 2, &devData);
    std::unordered_map<int, double> stats;
    for (int i = 0; i < len; ++i) {
        stats[devData[i].ddrNumaId] += devData[i].count / 1024 / 1024;
    }
    for (const auto &entry : stats) {
        int id = entry.first;
        double sum = entry.second;
        numaTotalDDRC[id] = sum;
    }
    numaNum = numaTotalDDRC.size();
    DevDataFree(devData);
    PmuDataFree(oriData);
    PmuDisable(pd);
}

// get numaId --> cpu core ids
void initNumaToCoreList()
{
    unsigned *coreList;
    for (unsigned i = 0; i < numaNum; ++i) {
        coreList = nullptr;
        int len = PmuGetNumaCore(i, &coreList);
        numaToCpuCore[i] = coreList;
        numaToCpuNumber[i] = len;
    }
}

// parse the CPU core range in the format "0-255" or "0-3,5"
std::vector<int> parseCpuRange(const std::string &rangeStr)
{
    std::vector<int> cpus;
    std::stringstream ss(rangeStr);
    std::string part;

    while(getline(ss, part, ',')) {
        size_t hyphen_pos = part.find("-");
        if (hyphen_pos != std::string::npos) {
            int start = std::stoi(part.substr(0, hyphen_pos));
            int end = std::stoi(part.substr(hyphen_pos + 1));
            if (start > end) {
                std::cerr << "Invalid CPU range: " << part << std::endl;
            }
            for (int i = start; i <= end; ++i) {
                cpus.push_back(i);
            }
        } else {
            cpus.push_back(std::stoi(part));
        }
    }

    std::sort(cpus.begin(), cpus.end());
    cpus.erase(unique(cpus.begin(), cpus.end()), cpus.end());
    return cpus;
}

// get cpu core of pid from /proc/[pid]/status
std::string getCpuAffinityList(int pid)
{
    std::string path = "/proc/" + std::to_string(pid) + "/status";
    std::ifstream in(path);
    if (!in.is_open()) {
        std::cerr << "Not found: " << path << std::endl;
        return "";
    }
    std::string line;
    const std::string targetKey = "Cpus_allowed_list:";
    while (getline(in, line)) {
        if (line.find(targetKey) == 0) {
            size_t pos = line.find("\t");
            if (pos == std::string::npos)
                pos = targetKey.length();
            return line.substr(pos + 1);
        }
    }
    return "";
}

int getCpuCore(int pid)
{
    try {
        std::string rangeStr = getCpuAffinityList(pid);
        if (rangeStr == "") {
            return -1;
        }
        pidBoundCpus = parseCpuRange(rangeStr);
    } catch (const std::exception &e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }
    return 0;
}

bool hasCommonCpu(const unsigned *cpuArray, size_t arraySize, const std::vector<int> &cpuVector)
{
    if (cpuArray == nullptr || arraySize == 0 || cpuVector.empty()) {
        return false;
    }

    if (arraySize < cpuVector.size()) {
        std::unordered_set<unsigned> arraySet(cpuArray, cpuArray + arraySize);
        for (const auto &cpu : cpuVector) {
            if (arraySet.count(cpu) > 0) {
                return true;
            }
        }
    } else {
        std::unordered_set<unsigned> vecSet(cpuVector.begin(), cpuVector.end());
        for (size_t i = 0; i < arraySize; ++i) {
            if (vecSet.count(cpuArray[i]) > 0) {
                return true;
            }
        }
    }

    return false;
}

std::string GetL3CMissPercent(unsigned llc_miss, unsigned llc_cache)
{
    std::ostringstream oss;
    double ratio = llc_cache != 0 ? static_cast<double>(llc_miss) / llc_cache * 100.0 : 0.0;
    oss << std::fixed << std::setprecision(FLOAT_PRECISION) << ratio;
    return oss.str();
}

// Find all processes containing the specified keyword
std::vector<int> findProcessesByKeyword(const std::string& keyword) {
    std::vector<int> matchedPids;
    DIR *dir = opendir("/proc");
    if (!dir) {
        std::cerr << "Failed to open /proc directory" << std::endl;
        return matchedPids;
    }

    struct dirent *entry;
    while ((entry = readdir(dir)) != nullptr) {
        // Check if directory name is a number (PID)
        if (entry->d_type == DT_DIR) {
            char *endptr;
            long pid = strtol(entry->d_name, &endptr, 10);
            if (*endptr == '\0') {  // Ensure the entire string is a number
                // Read process name
                std::string cmdlinePath = "/proc/" + std::string(entry->d_name) + "/cmdline";
                std::ifstream cmdlineFile(cmdlinePath);
                if (cmdlineFile.is_open()) {
                    std::string cmdline;
                    std::getline(cmdlineFile, cmdline);
                    // Check if command line contains keyword
                    if (cmdline.find(keyword) != std::string::npos) {
                        matchedPids.push_back(pid);
                    }
                }
            }
        }
    }
    closedir(dir);

    std::cout << "Found " << matchedPids.size() << " processes containing keyword '" << keyword << "'" << std::endl;
    for (const auto& pid : matchedPids) {
        std::cout << "PID: " << pid << std::endl;
    }

    return matchedPids;
}

// Collect LLC miss data for all threads of the processes (固定10秒统计时间)
std::vector<ThreadMissRatio> collectThreadLlcMissData(const std::vector<int>& processList) {
    std::vector<ThreadMissRatio> threads;

    if (processList.empty()) {
        return threads;
    }

    // Get CPU affinity for all processes
    std::vector<int> validPids;
    std::vector<int> allBoundCpus;
    std::unordered_set<int> uniqueCpus;

    for (const auto& pid : processList) {
        if (getCpuCore(pid) != -1) {
            validPids.push_back(pid);
            // Add each process's bound CPUs to the total list
            for (const auto& cpu : pidBoundCpus) {
                if (uniqueCpus.find(cpu) == uniqueCpus.end()) {
                    allBoundCpus.push_back(cpu);
                    uniqueCpus.insert(cpu);
                }
            }
        }
    }

    if (validPids.empty()) {
        return threads;
    }

    char* evtList[2];
    evtList[0] = (char*)"r33";  // LLC miss event
    evtList[1] = (char*)"r32";  // LLC cache access event
    PmuAttr attr = {0};
    attr.evtList = evtList;
    attr.numEvt = 2;
    attr.pidList = validPids.data();  // Pass all valid process PIDs
    attr.numPid = validPids.size();   // Set process count
    attr.cpuList = allBoundCpus.data();
    attr.numCpu = allBoundCpus.size();

    int pd = PmuOpen(COUNTING, &attr);
    if (pd == -1) {
        std::cerr << "PmuOpen failed for multiple PIDs" << std::endl;
        std::cerr << "Error msg: " << Perror() << std::endl;
        return threads;
    }

    PmuEnable(pd);

    // 固定统计10秒
    std::cout << "Collecting LLC miss data for " << duration << " seconds..." << std::endl;
    sleep(duration);

    PmuData* pmuData = nullptr;
    int len = PmuRead(pd, &pmuData);
    if (len == -1) {
        std::cerr << "PmuRead failed" << std::endl;
        std::cerr << "Error msg: " << Perror() << std::endl;
        PmuDisable(pd);
        PmuClose(pd);
        return threads;
    }

    // Create a mapping to store LLC miss and cache data for each thread
    std::unordered_map<unsigned, std::pair<unsigned, unsigned>> tidData; // tid --> (llc_miss, llc_cache)
    std::unordered_map<unsigned, int> tidToPid; // Record which process each thread belongs to
    std::unordered_map<unsigned, std::string> tidToComm; // Record process name for each thread

    // Process collected data, distinguish different threads by tid
    for (int i = 0; i < len; ++i) {
        PmuData& data = pmuData[i];
        unsigned tid = data.tid;
        int pid = data.pid;

        // Record which process the thread belongs to
        tidToPid[tid] = pid;

        // Record process name
        if (data.comm != nullptr) {
            tidToComm[tid] = std::string(data.comm);
        }

        if (strcmp(data.evt, "r33") == 0) {
            tidData[tid].first += data.count;  // llc_miss
        }
        if (strcmp(data.evt, "r32") == 0) {
            tidData[tid].second += data.count;  // llc_cache
        }
    }

    PmuDataFree(pmuData);
    PmuDisable(pd);
    PmuClose(pd);

    // Convert collected data to ThreadMissRatio objects
    for (const auto& entry : tidData) {
        unsigned tid = entry.first;
        unsigned llc_miss = entry.second.first;
        unsigned llc_cache = entry.second.second;
        int pid = tidToPid[tid];
        std::string comm = tidToComm.count(tid) ? tidToComm[tid] : "";
        threads.emplace_back(pid, tid, llc_miss, llc_cache, comm);
    }

    return threads;
}

// Print thread LLC miss ratio information
void PrintThreadMissRatio(const std::vector<ThreadMissRatio>& threads) {
    std::cout << std::string(120, '=') << std::endl;
    std::cout << "Thread LLC Miss Ratio Statistics (10s Collection)" << std::endl;
    std::cout << std::string(120, '-') << std::endl;
    std::cout << std::setw(6) << "Index" << std::setw(8) << std::left << "PID"
              << std::setw(8) << "TID" << std::setw(16) << "Process"
              << std::setw(10) << "LLC Miss%" << std::setw(12) << "Miss Count"
              << std::setw(12) << "Cache Count" << std::setw(8) << "Active" << std::endl;
    std::cout << std::string(120, '-') << std::endl;

    for (size_t i = 0; i < threads.size(); ++i) {
        std::string process_name = threads[i].comm.empty() ? "unknown" : threads[i].comm;
        // Truncate process name if too long
        if (process_name.length() > 14) {
            process_name = process_name.substr(0, 11) + "...";
        }
        
        std::cout << std::setw(6) << i << std::setw(8) << std::left << threads[i].pid
                  << std::setw(8) << threads[i].tid << std::setw(16) << process_name
                  << std::setw(10) << GetL3CMissPercent(threads[i].llc_miss, threads[i].llc_cache) + "%"
                  << std::setw(12) << threads[i].llc_miss
                  << std::setw(12) << threads[i].llc_cache
                  << std::setw(8) << (threads[i].is_active ? "Yes" : "No") << std::endl;
    }

    std::cout << std::string(120, '_') << std::endl;
    
    // 统计活跃线程数量
    int active_count = 0;
    for (const auto& thread : threads) {
        if (thread.is_active) {
            active_count++;
        }
    }
    std::cout << "Total threads: " << threads.size() << ", Active threads: " << active_count 
              << " (cache access >= " << CACHE_ACCESS_THRESHOLD << ")" << std::endl;
}

// 优化的线程绑定策略：只绑定活跃线程，按LLC miss率配对
void optimizedBindThreads(std::vector<ThreadMissRatio>& threads) {
    if (threads.empty()) {
        std::cout << "No valid thread data found" << std::endl;
        return;
    }

    // 过滤出活跃线程
    std::vector<ThreadMissRatio> activeThreads;
    for (const auto& thread : threads) {
        if (thread.is_active) {
            activeThreads.push_back(thread);
        }
    }

    if (activeThreads.empty()) {
        std::cout << "No active threads found (all threads have cache access < " << CACHE_ACCESS_THRESHOLD << ")" << std::endl;
        return;
    }

    // 按LLC miss率排序（降序）
    std::sort(activeThreads.begin(), activeThreads.end());

    // 获取物理核心映射
    std::vector<std::pair<int, int>> physicalCores = getPhysicalCoreMapping();
    if (physicalCores.empty()) {
        std::cerr << "Failed to get physical core mapping" << std::endl;
        return;
    }

    std::cout << std::string(80, '=') << std::endl;
    std::cout << "Optimized Thread Binding Strategy" << std::endl;
    std::cout << "Active threads: " << activeThreads.size() << " / " << threads.size() << std::endl;
    std::cout << "Available physical cores: " << physicalCores.size() << std::endl;
    std::cout << std::string(80, '-') << std::endl;

    // 配对绑定：高miss率线程与低miss率线程绑定到同一物理核心
    size_t maxBindings = std::min(activeThreads.size() / 2, physicalCores.size());
    
    for (size_t i = 0; i < maxBindings; ++i) {
        // 高miss率线程（前半部分）
        ThreadMissRatio& highMissThread = activeThreads[i];
        // 低miss率线程（后半部分）
        ThreadMissRatio& lowMissThread = activeThreads[activeThreads.size() - 1 - i];
        
        // 绑定到同一物理核心的两个逻辑核心
        int core0 = physicalCores[i].first;   // 偶数核心
        int core1 = physicalCores[i].second;  // 奇数核心
        
        // 绑定高miss率线程到偶数核心
        bool success1 = bindThreadToCore(highMissThread.tid, core0);
        // 绑定低miss率线程到奇数核心
        bool success2 = bindThreadToCore(lowMissThread.tid, core1);
        
        if (success1 && success2) {
            std::string high_process = highMissThread.comm.empty() ? "unknown" : highMissThread.comm;
            std::string low_process = lowMissThread.comm.empty() ? "unknown" : lowMissThread.comm;
            
            if (high_process.length() > 12) high_process = high_process.substr(0, 9) + "...";
            if (low_process.length() > 12) low_process = low_process.substr(0, 9) + "...";
            
            std::cout << "Physical Core " << i << " (Cores " << core0 << "," << core1 << "):" << std::endl;
            std::cout << "  High Miss: TID " << std::setw(6) << highMissThread.tid 
                      << " (" << std::setw(12) << std::left << high_process << ") "
                      << "LLC Miss: " << std::setw(8) << GetL3CMissPercent(highMissThread.llc_miss, highMissThread.llc_cache) + "% "
                      << "-> Core " << core0 << std::endl;
            std::cout << "  Low Miss:  TID " << std::setw(6) << lowMissThread.tid 
                      << " (" << std::setw(12) << std::left << low_process << ") "
                      << "LLC Miss: " << std::setw(8) << GetL3CMissPercent(lowMissThread.llc_miss, lowMissThread.llc_cache) + "% "
                      << "-> Core " << core1 << std::endl;
        }
    }

    std::cout << std::string(80, '_') << std::endl;
    std::cout << "Successfully bound " << maxBindings * 2 << " threads to " << maxBindings << " physical cores" << std::endl;
    
    // 显示未绑定的活跃线程
    if (activeThreads.size() > maxBindings * 2) {
        std::cout << "Unbound active threads: " << (activeThreads.size() - maxBindings * 2) << std::endl;
        for (size_t i = maxBindings * 2; i < activeThreads.size(); ++i) {
            std::string process_name = activeThreads[i].comm.empty() ? "unknown" : activeThreads[i].comm;
            if (process_name.length() > 12) process_name = process_name.substr(0, 9) + "...";
            std::cout << "  TID " << std::setw(6) << activeThreads[i].tid 
                      << " (" << std::setw(12) << std::left << process_name << ") "
                      << "LLC Miss: " << std::setw(8) << GetL3CMissPercent(activeThreads[i].llc_miss, activeThreads[i].llc_cache) + "%" << std::endl;
        }
    }
}

void print_usage() {
    std::cerr << "Usage: llc_miss_thread_bind <threshold> <topNum> <period> <process_keyword>\n";
    std::cerr << "--threshold : the collect threshold of total ddrc bandwidth, unit M/s\n";
    std::cerr << "--topNum : the top N thread of llc miss ratio collection\n";
    std::cerr << "--period : the period of llc_miss_ratio collect, unit ms\n";
    std::cerr << "--process_keyword : the keyword to match process names\n";
    std::cerr << " example: llc_miss_thread_bind 100 10 1000 stress-ng\n";
    std::cerr << "\n";
    std::cerr << "Features:\n";
    std::cerr << "- Fixed 10-second LLC miss ratio collection\n";
    std::cerr << "- Activity filtering: only bind threads with cache access >= " << CACHE_ACCESS_THRESHOLD << "\n";
    std::cerr << "- Complementary binding: high miss rate + low miss rate threads on same physical core\n";
    std::cerr << "- Process-level LLC miss ratio collection and binding\n";
}

int main(int argc, char** argv)
{
    if (argc < 5) {
        print_usage();
        return 0;
    }
    double threshold = 0.0;
    std::string processKeyword;
    bool collectL3CMissFlag = false;

    try {
        threshold = std::stod(argv[1]);
        if (threshold <= 0) {
            throw std::invalid_argument("threshold must be a positive number.");
        }

        topNum = std::stod(argv[2]);
        if (topNum <= 0) {
            throw std::invalid_argument("TopNum must be a positive number.");
        }

        period = std::stoi(argv[3]);
        if (period <= 0) {
            throw std::invalid_argument("Period must be a positive integer.");
        }

        processKeyword = argv[4];
    } catch (const std::exception& e) {
        std::cerr << "Error parsing arguments: " << e.what() << "\n";
        print_usage();
        return EXIT_FAILURE;
    }

    // Find all processes containing the keyword
    std::vector<int> matchedPids = findProcessesByKeyword(processKeyword);
    if (matchedPids.empty()) {
        std::cerr << "No processes found containing keyword '" << processKeyword << "'" << std::endl;
        return EXIT_FAILURE;
    }

    totalDDRCBandwidth();
    initNumaToCoreList();

    // Check if NUMA bandwidth exceeds threshold
    for (const auto &data : numaTotalDDRC) {
        std::cout << "Numa ID: " << data.first << ", total bandwidth: " << data.second << "M/s";
        if (data.second > threshold) {
            std::cout << " --> Exceeds threshold, need to collect LLC miss ratio";
            collectL3CMissFlag = true;
        } else {
            std::cout << " --> Does not exceed threshold";
        }
        std::cout << std::endl;
    }

    if (collectL3CMissFlag) {
        // Collect LLC miss ratio for all threads of matching processes (固定10秒)
        std::vector<ThreadMissRatio> threads = collectThreadLlcMissData(matchedPids);

        // Print thread LLC miss ratio information
        PrintThreadMissRatio(threads);

        // 使用优化的绑定策略
        optimizedBindThreads(threads);
    }

    return 0;
}